import brian as brian
from brian import mV, ms, second

    
class SynfireNetwork:
           
    # Neuron model parameters
    Vr = -70*mV
    Vt = -55*mV
    taum = 10*ms
    taupsp = 0.325*ms
    weight = 4.86 * mV
    refractoriness = 1 * ms
    
    # binSize used for packet detection
    binSize = 0.2 * ms 
    Input = None
    SpikeMonitors = None
    
    # class constructor
    def __init__(self, n_layers=10, neurons_per_layer=100):
        # create the network (without any inputs)

        self.n_layers = n_layers
        self.neurons_per_layer = neurons_per_layer

        self.StateMonitors = []
        self.SetupNetwork()
        
    
    # Creates the neurons and their connections, also adds monitors
    def SetupNetwork(self):
        # Neuron model
        eqs = brian.Equations('''
                dV/dt=(-(V-self.Vr)+x)*(1./self.taum) : volt
                dx/dt=(-x+y)*(1./self.taupsp) : volt
                dy/dt=-y*(1./self.taupsp)+25.27*mV/ms+\
                (39.24*mV/ms**0.5)*xi : volt
                ''')
        
        # Neuron groups
        self.Neurons = brian.NeuronGroup(N=self.neurons_per_layer*self.n_layers, 
                model=eqs,
                threshold=self.Vt,reset=self.Vr,refractory= self.refractoriness)
        
        # The network structure
        self.Layers = [self.Neurons.subgroup(self.neurons_per_layer)
                for i in range(self.n_layers)]
        self.NeuronConnection = brian.DelayConnection(self.Neurons, self.Neurons,'y', max_delay = 2 * ms)
        
        for i in range(self.n_layers-1):
            self.NeuronConnection.connect_full(self.Layers[i], self.Layers[i+1],self.weight, delay = (0 * ms, 2 * ms))
            
        self.Network = brian.Network(self.Neurons, self.NeuronConnection)
        
        self.AddMonitors()
    
    # Extracts the list of spiketimes from the spike monitor
    def GetSpikeTimes(self, spikeMonitor):
        return [spiketime for spiketimes in spikeMonitor.spiketimes.itervalues() for spiketime in spiketimes]
    
    # Calculates the average spike count per  bin (10% are added)
    def CalcAverageSpikeCountPerBin(self):
        self.AverageSpikeCountsPerBin = []
        
        # Create bins and count number of spikes in each bin
        for SpikeMonitor in self.SpikeMonitors:
           averageSpikeCountPerBinOfLayer = SpikeMonitor.nspikes / (self.lastRunTime / self.binSize) * 1.1
           if averageSpikeCountPerBinOfLayer < 2:
               averageSpikeCountPerBinOfLayer = 2
               
           self.AverageSpikeCountsPerBin.append(averageSpikeCountPerBinOfLayer)
                            
      
    # Connects the given input with the first layer of the network
    def SetInput(self, input):
        self.Input = input
        self.InputConnection = brian.Connection(self.Input, self.Layers[0], 'y')
        self.InputConnection.connect_full(weight=self.weight)
        
        self.SpikeMonitors = [brian.SpikeMonitor(self.Input)] + self.SpikeMonitors
        
        self.Network.add(self.Input, self.InputConnection, self.SpikeMonitors)
        
    # Adds all monitors required to the network
    def AddMonitors(self):
        # Setup monitors
        
        # Record state variables like standard deviation
        self.StateMonitor =  [brian.StateMonitor(p, 'y') for p in self.Layers]
        
        # Record the single spikes
        self.SpikeMonitors = [brian.SpikeMonitor(p) for p in self.Layers]
        
        # Record population rates
        self.PSTHMonitor = [brian.PopulationRateMonitor(source,0.2*ms) for source in self.Layers]
        
        self.Network.add(self.StateMonitor, self.SpikeMonitors, self.PSTHMonitor)
    
    def AddStateMonitor(self, layer, state='V'):
        """Add a new monitor which records an arbitrary neuronal state
        variable (such as membrane potential, psp, etc.)
        """

        state_monitor = brian.StateMonitor(self.Layers[layer],
                                            state, record=True)
        self.StateMonitors.append(state_monitor)
        self.Network.add(state_monitor)
        
    # Runs a simulation
    def Run(self, trialTime):
            
        # Setup the network, and run it
        self.Neurons.V = self.Vr + brian.rand(len(self.Neurons)) * (self.Vt-self.Vr)
        
        self.lastRunTime = trialTime
        
        self.Network.run(trialTime)
        #self.Network.reinit()
        
        self.CalcAverageSpikeCountPerBin()

        self.FindPackets()
    
    # Finds a packet within each layer, if available
    def FindPackets(self):
        def GetSpikeCountInBin(SpikeTimes, currentPosition):
            return len(filter(lambda x: x * second >= currentPosition and x * second  < currentPosition + self.binSize, SpikeTimes))
        
        self.Packets = []
        
        for SpikeMonitor in self.SpikeMonitors:
            SpikeTimes = self.GetSpikeTimes(SpikeMonitor)
            averageSpikeCountPerBin = self.AverageSpikeCountsPerBin[self.SpikeMonitors.index(SpikeMonitor) - 1] 
            
            currentPosition = 0 * ms
            
            packetStart = -1 * ms
            packetEnd   = -1 * ms
            searchEnd   = False
            
            while currentPosition < self.lastRunTime:
                if GetSpikeCountInBin(SpikeTimes, currentPosition) > averageSpikeCountPerBin:
                    if packetStart == -1 * ms:
                        packetStart = currentPosition
                    
                    searchEnd = True
                else:
                    if searchEnd:
                        packetEnd = currentPosition
                        searchEnd = False
                    
                   
                currentPosition += self.binSize
            
            if packetStart != -1 * ms:
                if searchEnd:
                    packetEnd = self.lastRunTime
                 
                self.Packets.append((packetStart, packetEnd))
            else:
                self.Packets.append(None)
    
    # Returns the standard deviation of all spike within a packet detected
    # or the standard deviation of all spikes if no packet has been found
    def GetClippedStandardDeviation(self, SpikeMonitor, packetPosition):
        SpikeTimes = self.GetSpikeTimes(SpikeMonitor)
        
        if packetPosition is None:
            return brian.std(SpikeTimes)
        else:
            return brian.std(filter(lambda x: x * second >= packetPosition[0]  and x * second <= packetPosition[1], SpikeTimes))
    
    # Plots the results of the last test run
    def ShowResults(self):
        
        # Plot results
        
        # Plot standard deviation by layer
        brian.subplot(321)
        brian.plot(map(lambda x: self.GetClippedStandardDeviation(x, self.Packets[self.SpikeMonitors.index(x)]), self.SpikeMonitors), 'b-', linewidth=2)
        brian.xlabel('Layer')
        brian.ylabel(r'$\sigma$')
        
        # Plot spike count per layer
        brian.subplot(322)
        brian.plot(map(lambda x: x.nspikes, self.SpikeMonitors), 'b-', linewidth=2)
        brian.xlabel('Layer')
        brian.ylabel('n Spikes')
        
        # Plot spikes by time
        brian.subplot(323)
        brian.raster_plot(showgrouplines=True,*self.SpikeMonitors)
        
        for packet in self.Packets:
            if not packet is None:
                index = self.Packets.index(packet)
                brian.plot([packet[0] / ms, packet[0] / ms], [index, index + 1], 'g--', linewidth=2)
                brian.plot([packet[1] / ms, packet[1] / ms], [index, index + 1], 'r--', linewidth=2)
        
        
        # Plot histogram
        brian.subplot(324)
        for m in self.PSTHMonitor:
            brian.plot(m.times/ms, m.rate)
        
            
        brian.xlabel("time (ms)")
        brian.xlabel("population rate (Hz)")

        for i, monitor in enumerate(self.StateMonitors):
            brian.subplot(3,2,5+i)
            for neuron_state in monitor:
                 brian.plot(monitor.times/ms, neuron_state)
                 brian.xlabel("time (ms)")
                 brian.ylabel(monitor.varname)
            
        
        brian.show()
        
    
model = SynfireNetwork(n_layers=4)
model.weight = 100*mV
model.SetInput(brian.PulsePacket(t=10*ms,n=50,sigma=0*ms))
model.AddStateMonitor(0)
model.Run(20 * ms)
model.ShowResults()





    



